import re
import matplotlib.pyplot as plt
import argparse

def extract_loss_from_log(log_file_path):
    """
    从日志文件中提取 loss 值。
    
    参数:
        log_file_path (str): 日志文件的路径。
    
    返回:
        list: 包含所有提取到的 loss 值的列表。
    """
    loss_values = []
    with open(log_file_path, 'r') as file:
        for line in file:
            # 使用正则表达式匹配 loss 值
            match = re.search(r'loss:([\d.]+)', line)
            if match:
                # 将匹配到的 loss 值转换为浮点数并添加到列表中
                loss_value = float(match.group(1))
                loss_values.append(loss_value)
    return loss_values

def plot_loss(loss_values, save_path):
    """
    绘制 loss 值的变化曲线图，并保存图像。
    
    参数:
        loss_values (list): 包含 loss 值的列表。
        save_path (str): 保存图像的路径。
    """
    # 生成 x 轴的步数（假设每一步对应一个 loss 值）
    steps = list(range(len(loss_values)))
    
    # 绘制曲线图
    plt.figure(figsize=(10, 6))
    plt.plot(steps, loss_values, label=args.save_name, color='orange', linestyle='-', marker='o', markersize=3, alpha=0.7)
    plt.title('loss')
    plt.xlabel('Step')
    plt.ylabel('Loss')
    plt.legend(loc='upper right', fontsize='small')  # 设置图例位置和字体大小
    plt.grid(True, linestyle='--', alpha=0.5)  # 设置网格线样式
    plt.tight_layout()  # 自动调整子图参数，使之填充整个图像区域
    
    # 保存图像
    plt.savefig(save_path)
    print(f"图像已保存到: {save_path}")
    plt.show()
    
if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Chat with MiniMind")
    parser.add_argument('--log_file_path', default='None', type=str)
    parser.add_argument('--save_name', default='out', type=str)
    parser.add_argument('--save_path', default='out', type=str)
    args = parser.parse_args()

    # 图像保存路径
    save_path = args.save_path + '/' + args.save_name + '.png'  # 替换为你希望保存图像的路径
    # 提取 loss 值
    loss_values = extract_loss_from_log(args.log_file_path)
    
    # 绘制曲线图并保存
    plot_loss(loss_values, save_path)